(* $Id: xCoreSyntaxExt.ml 221 2009-03-06 16:24:38Z yann.regisgianas $ *)

open XCoreSyntax

let atom_id a =
  Var.Atom.basename a

let atom_position a =
  PIdentifier.position (Var.Atom.basename a)

let atom_as_string a =
  PIdentifier.as_string (Var.Atom.basename a)

let mk_lformula pos f =
  { fpos = pos; fvalue = f }

let rec remove_fannot' f =
  match f.fvalue with
    | FAnnot (f, ty) -> remove_fannot' f
    | _ -> f

let remove_fannot = function
  | FAnnot (t, _) -> t.fvalue
  | t -> t

let mk_fannot pos f ty =
  FAnnot (mk_lformula pos f, ty)

let explicit_formula = function
  | ExplicitFormula f -> f
  | _ -> assert false

let mk_formula_type_scheme vs ty =
  FTScheme (create_formula_scheme_abs (vs, ty))

let mk_flam pos bs f =
  FLam (create_lfun_abs (bs, mk_lformula pos f))

let mk_fprod pos fs =
  FProd (List.map (mk_lformula pos) fs)

let rec destruct_fprod = function
  | FProd [ t ] -> destruct_fprod t.fvalue
  | FProd ts -> ts
  | FAnnot (t, _) -> destruct_fprod t.fvalue
  | t -> [ mk_lformula Position.dummy t ]

let mk_fprod_type = function
  | [ ty ] -> ty
  | tys -> FTProd tys

let mk_conj pos f1 f2 =
  FApp (mk_lformula pos (FPrimitive PAnd),
	List.map (mk_lformula pos) [ f1; f2 ])

let rec mk_conjs pos = function
  | [] -> FTrue
  | [ x ] -> x
  | [ p1; p2 ] -> mk_conj pos p1 p2
  | x :: xs -> mk_conj pos x (mk_conjs pos xs)

let rec destruct_conjs g =
  match (remove_fannot' g).fvalue with
  | FApp (f, args) ->
      (match (remove_fannot' f).fvalue, args with
	 | FPrimitive PAnd, [ t1; t2 ] ->
	     destruct_lconjs t1 @ destruct_lconjs t2
	 | _ -> [g])
  | _ -> [g]

and destruct_lconjs f =
  destruct_conjs f

let mk_disj pos f1 f2 =
  FApp (mk_lformula pos (FPrimitive POr),
	List.map (mk_lformula pos) [ f1; f2 ])

let rec mk_disjs pos = function
  | [] -> FTrue
  | [ x ] -> x
  | [ p1; p2 ] -> mk_disj pos p1 p2
  | x :: xs -> mk_disj pos x (mk_disjs pos xs)

let mk_fimply pos f1 f2 =
  FApp (mk_lformula pos (FPrimitive PImply),
	List.map (mk_lformula pos) [ f1; f2 ])

let mk_fequiv pos f1 f2 =
  FApp (mk_lformula pos (FPrimitive PEquiv),
	List.map (mk_lformula pos) [ f1; f2 ])

let mk_fforalltys pos vs x =
  FForallTys (create_fforalltys_abs (vs, mk_lformula pos x))

let mk_lforall pos bs x =
  FForall (create_lforall_abs (bs, [], mk_lformula pos x))

let mk_lforall_triggers pos bs ts x =
  FForall (create_lforall_abs (bs, ts, mk_lformula pos x))

let mk_lexists pos bs x =
  FExists (create_lexists_abs (bs, mk_lformula pos x))

let mk_fapp pos f args =
  FApp (mk_lformula pos f, List.map (mk_lformula pos) args)

let rec destruct_fapp f =
  match f.fvalue with
    | FApp (f, args) ->
	let (f, argss) = destruct_fapp f in
	  (remove_fannot f), args :: argss
    | FAnnot (f, _) ->
	destruct_fapp f
    | x ->
	x, []

let destruct_fapp f =
  let (f, argss) = destruct_fapp f in
    (f, List.rev argss)

let mk_fkapp pos f args =
  FKApp (f, List.map (mk_lformula pos) args)

let mk_feq pos ty v1 v2 =
  FEq (ty, mk_lformula pos v1, mk_lformula pos v2)

let mk_fneq pos ty v1 v2 =
  FApp (mk_lformula pos (FPrimitive PNot),
	[ mk_lformula pos (mk_feq pos ty v1 v2) ])

let mk_conj pos f1 f2 =
  mk_fapp pos (FPrimitive PAnd) [ f1; f2 ]

let mk_type_scheme vs ty =
  TScheme (create_scheme_abs (vs, ty))

let mk_prod_type = function
  | [ ty ] -> ty
  | ts -> TProd ts

let monoscheme ty =
  TScheme (create_scheme_abs ([], ty))

let mk_lterm pos t =
  { tpos = pos; tvalue = t }

let mk_prod = function
  | [ t ] -> t.tvalue
  | ts -> EProd ts

let mk_lprod pos ts =
  mk_prod (List.map (mk_lterm pos) ts)

let as_eid = function
  | EId x -> x
  | _ -> assert false

let full_destruct_let_bindings bs =
  let (xs, vvs, tys, bs) =
    Misc.list_map_to_4_lists
      (fun (x, TScheme abs) ->
	 let (vs, ty) = open_scheme_abs abs in
	   (x, vs, ty, (x, ty)))
      bs
  in
    (xs, List.flatten vvs, tys, bs)

class map_with_pos = object (self)
  inherit map

  val mutable pos = Position.dummy

  method current_position =
    pos

  method lformula f =
    pos <- f.fpos;
    { fpos = f.fpos; fvalue = self#formula f.fvalue }

  method lterm t =
    pos <- t.tpos;
    { tpos = t.tpos; tvalue = self#term t.tvalue }

end

class ['a] fold_with_pos = object (self)
  inherit ['a] fold

  val mutable pos = Position.dummy

  method current_position =
    pos

  method lformula acu f =
    pos <- f.fpos;
    self#formula acu f.fvalue

  method lterm acu t =
    pos <- t.tpos;
    self#term acu t.tvalue

end

class foralltys_extruder =
object (self)

  inherit map

  val mutable rs = Var.AtomSet.empty

  method rigid_vars = Var.AtomSet.fold (fun s x -> s :: x) rs []

  method fforalltys abs =
    let (vs, t) = open_fforalltys_abs abs in
      rs <- List.fold_left (fun s x -> Var.AtomSet.add x s)  rs vs;
      self#formula t.fvalue

end

let extrude_foralltys f =
  let extruder = new foralltys_extruder in
  let f = extruder#lformula f in
    mk_lformula f.fpos
      (FForallTys (create_fforalltys_abs (extruder#rigid_vars, f)))

class substituer phi =
object

  inherit map

  method fvar y =
    try
      snd (List.find (fun (z, t) ->
                    (*                      Printf.eprintf "%s =?= %s => %B \n%!"
                          (PIdentifier.as_string (Var.Atom.basename z))
                          (PIdentifier.as_string (Var.Atom.basename y))
                          (Var.Atom.equal z y); *)
                        Var.Atom.equal z y) phi)
    with Not_found ->
      FVar y

end

let subst_formula (x, v) f =
  extrude_foralltys ((new substituer [(x, v)])#lformula f)

let many_subst_formula phi f =
  extrude_foralltys ((new substituer phi)#lformula f)

let as_fvar = function
  | FVar f -> f
  | _ -> assert false

exception NotInductiveCaseSyntax

let rec destruct_inductive_case = function
  | FForallTys abs ->
      let (vs, t) = open_fforalltys_abs abs in
      let (ts, bs, hypotheses, conclusion) =
	destruct_inductive_case t.fvalue
      in
	(vs @ ts, bs, hypotheses, conclusion)

  | FForall abs ->
      let (ins, _, t) = open_lforall_abs abs in
      let (ts, bs, hypotheses, conclusion) =
	destruct_inductive_case t.fvalue
      in
	(ts, ins @ bs, hypotheses, conclusion)

  | ((FApp (g, args)) as f) ->
      (match remove_fannot g.fvalue, args with
	 | FPrimitive PImply, [ h; c]  ->
	     ([], [], [h], c.fvalue)

	 | _ ->
	     ([], [], [], f))

  | FAnnot (f, _) ->
      destruct_inductive_case f.fvalue

  | _ ->
      raise NotInductiveCaseSyntax

let destruct_implication f =
  try
    let pos = f.fpos in
    let (ts, bs, hs, conclusion) = destruct_inductive_case f.fvalue in
      if hs = [] then
	[ f ]
      else
	let hs = List.map (fun f -> f.fvalue) hs in
	let cs = destruct_conjs (mk_lformula pos conclusion) in
	let cs = List.map (fun f -> f.fvalue) cs in
	  List.map (fun c ->
		      mk_lformula pos
			(mk_fforalltys pos ts
			   (mk_lforall pos bs
			      (mk_fimply pos (mk_conjs pos hs) c))))
	    cs
  with NotInductiveCaseSyntax -> [ f ]

let rec destruct_poly_flam = function
  | FForallTys abs ->
      let (vs, t) = open_fforalltys_abs abs in
      let (vs', ins, body) = destruct_poly_flam t.fvalue in
	(vs @ vs', ins, body)

  | FLam abs ->
      let (ins, body) = open_lfun_abs abs in
      let (vs, ins', body) = destruct_poly_flam body.fvalue in
	(vs, ins @ ins', body)

  | FAnnot (t, ty) ->
      destruct_poly_flam t.fvalue

  | t ->
      ([], [], t)

let rec destruct_poly_flam' = function
  | FForallTys abs ->
      let (vs, t) = open_fforalltys_abs abs in
      let (vs', ins, body) = destruct_poly_flam' t.fvalue in
	(vs @ vs', ins, body)

  | FLam abs ->
      let (ins, body) = open_lfun_abs abs in
	([], ins, body.fvalue)

  | t ->
      ([], [], t)

let destruct_formula_type_scheme_as_poly_bindings ?ids pos (FTScheme abs) =
  let (vs, ty) = open_formula_scheme_abs abs in
  let ids_from tys =
    match ids with
      | Some ids -> ids
      | None ->
	  List.map
	    (fun _ ->
	       Var.Atom.freshb (PIdentifier.fresh_value_id (pos, "x")))
	    tys
  in
    match ty with
      | FTCArrow (FTProd tys, ret_ty) ->
	    (vs, List.combine (ids_from tys) tys, ret_ty)

      | FTCArrow (ty, ret_ty) ->
	  (vs, List.combine (ids_from [ty]) [ ty ], ret_ty)

      | FTArrow (FTProd tys, ret_ty) ->
	  (vs, List.combine (ids_from tys) tys, ret_ty)

      | FTArrow (ty, ret_ty) ->
	  (vs, List.combine (ids_from [ty]) [ ty ], ret_ty)

      | _ -> (vs, [], ty)

let formula_monoscheme ty =
  FTScheme (create_formula_scheme_abs ([], ty))

let rec pattern_as_binding acu = function
  | PVar (x, ty) -> (x, ty) :: acu
  | PApp (_, ps) | PProd ps -> List.fold_left pattern_as_binding acu ps

let pattern_as_bindings = pattern_as_binding []

let rec pattern_as_logic_value pos = function
  | PVar (x, _) -> FVar x
  | PApp (k, ps) -> FKApp (k, List.map (pattern_as_logic_lvalue pos) ps)
  | PProd ps -> FProd (List.map (pattern_as_logic_lvalue pos) ps)

and pattern_as_logic_lvalue pos v =
  mk_lformula pos (pattern_as_logic_value pos v)

let eq_dataconstructor_app pos x (d, ty) =
  let s = mk_formula_type_scheme [] ty in
  let (_, bs, rty) = destruct_formula_type_scheme_as_poly_bindings pos s in
  let xs = List.map (fun (x, ty) -> mk_fannot pos (FVar x) ty) bs in
    mk_lexists pos bs
      (mk_feq pos rty
	 (mk_fannot pos (mk_fkapp pos d xs) rty)
	 (mk_fannot pos (FVar x) rty))

let neq_dataconstructor_app pos (d1, ty1) (d2, ty2) =
  let s1 = mk_formula_type_scheme [] ty1 in
  let (_, bs1, rty1) = destruct_formula_type_scheme_as_poly_bindings pos s1 in
  let xs1 = List.map (fun (x, ty) -> mk_fannot pos (FVar x) ty) bs1 in
  let s2 = mk_formula_type_scheme [] ty2 in
  let (_, bs2, rty2) = destruct_formula_type_scheme_as_poly_bindings pos s2 in
  let xs2 = List.map (fun (x, ty) -> mk_fannot pos (FVar x) ty) bs2 in
    mk_lforall pos (bs1 @ bs2)
      (mk_fneq pos rty1
	 (mk_fannot pos (mk_fkapp pos d1 xs1) rty1)
	 (mk_fannot pos (mk_fkapp pos d2 xs2) rty2))


let inhabited_types pos v ty =
  let x = Var.Atom.freshb (PIdentifier.fresh_value_id (pos, "x")) in
    mk_lexists pos [ (x, ty) ]
      (mk_feq pos ty
	 (mk_fannot pos v ty)
	 (mk_fannot pos (FVar x) ty))

let rec is_deferred_formula f =
  match f.fvalue with

    | FForallTys abs ->
	let (vs, f) = open_fforalltys_abs abs in
	  is_deferred_formula f

    | FAnnot (f, _) ->
	is_deferred_formula f

    | FDeferred ->
	true

    | f ->
	false

let arity_of_kind = function
  | KStar -> 0
  | KArrow ks -> List.length ks

let formula_monoscheme f =
  FTScheme (create_formula_scheme_abs ([], f))

let monoscheme ty =
  TScheme (create_scheme_abs ([], ty))

let as_pre_spec f =
  match (remove_fannot' f).fvalue with
    | FEq (_, lhs, rhs) ->
	(match destruct_fapp lhs with
	   | FPrimitive Pre, [ [fn] ] ->
	       (match remove_fannot fn.fvalue with
		  | FVar fn -> Some (fn, rhs)
		  | _ -> None)
	   | _ -> None)
    | _ -> None

let as_post_spec f =
  match remove_fannot f.fvalue with
    | FEq (_, lhs, rhs) ->
	(match destruct_fapp lhs with
	   | FPrimitive Post, [ [fn] ] ->
	       (match remove_fannot fn.fvalue with
		  | FVar fn -> Some (fn, rhs)
		  | _ -> None)
	   | _ -> None)
    | _ -> None


let rec destruct_poly pos = function
  | FForallTys abs ->
      let (ts, lf) = open_fforalltys_abs abs in
      let (ts', lf) = destruct_poly pos lf.fvalue in
        (ts @ ts', lf)
  | lf ->
      ([], mk_lformula pos lf)

let close_poly pos (ts, lf) =
  match ts with
    | [] -> lf
    | ts -> FForallTys (create_fforalltys_abs (ts, mk_lformula pos lf))

let as_function_spec f =
  let (ts, f) = destruct_poly f.fpos (remove_fannot' f).fvalue in
  match destruct_lconjs (remove_fannot' f) with
    | [ pre ; post ] ->
	(match as_pre_spec pre, as_post_spec post with
	   | Some (fn1, pre), Some (fn2, post) ->
	       if Var.Atom.equal fn1 fn2 then
		 Some (fn1, close_poly f.fpos (ts, FProd [ pre; post ]))
	       else None
	   | _ -> None)
    | l -> None
